import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules


class AssociationAnalyzer:
    def __init__(self, filename):
        self.filename = filename
        self.df = self.load_data()

    def load_data(self):
        # 加载数据
        file_path = "../InputData/"
        return pd.read_csv(file_path + self.filename + ".tsv", sep="\t")

    def map_effectiveness(self, value):
        # 处理effectiveness字段
        return (
            0
            if value in ["Considerably Effective", "Highly Effective"]
            else 1 if value in ["Moderately Effective", "Marginally Effective"] else 2
        )

    def map_side_effects(self, value):
        # 处理sideEffects字段
        return (
            2
            if value in ["Extremely Severe Side Effects", "Severe Side Effects"]
            else 1 if value == "Moderate Side Effects" else 0
        )

    def map_review_length(self, review):
        # 处理评论字段，计算字数并分类，忽略非字符串类型的数据
        if pd.isna(review) or not isinstance(review, str):
            return 0
        length = len(review)
        return (
            0
            if length <= 200
            else (
                1
                if 201 <= length <= 400
                else 2 if 401 <= length <= 600 else 3 if 601 <= length <= 800 else 4
            )
        )

    def ConditionDevide(self):
        # 打印列名以检查 'condition' 列是否存在
        print("Columns in DataFrame:", self.df.columns)

        # 对'condition'列进行分类，并统计每种'condition'的数量
        if "condition" in self.df.columns:
            condition_counts = self.df["condition"].value_counts()
            # 筛选出数量大于5的分类
            condition_counts_filtered = condition_counts[condition_counts > 5]
            # 使用seaborn绘制条形图
            plt.figure(figsize=(10, 8))  # 设置图表大小
            sns.barplot(
                x=condition_counts_filtered.index, y=condition_counts_filtered.values
            )
            plt.title("Condition Count Distribution (Counts > 5)")  # 设置图表标题
            plt.xlabel("Condition")  # 设置x轴标签
            plt.ylabel("Count")  # 设置y轴标签
            plt.xticks(rotation=90)  # 旋转x轴标签，以便更好地显示
            plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
            plt.show()
        else:
            print("Error: 'condition' column not found in the DataFrame.")

    def StarDevide(self):
        # 筛选出数量大于10的condition
        if "condition" in self.df.columns:
            condition_counts = self.df["condition"].value_counts()
            conditions_greater_than_10 = condition_counts[condition_counts > 10].index
            # 筛选出这些condition对应的数据
            df_filtered = self.df[self.df["condition"].isin(conditions_greater_than_10)]
            # 对每个condition的rating取平均值并向上取整
            average_ratings_ceil = (
                df_filtered.groupby("condition")["rating"].mean().apply(np.ceil)
            )
            # 创建一个足够大的画布
            n_conditions = len(average_ratings_ceil)
            n_rows = int(np.ceil(np.sqrt(n_conditions)))
            n_cols = int(np.ceil(n_conditions / n_rows))
            fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 6, n_rows * 6))
            # 为每个condition绘制饼状图
            for i, (condition, avg_rating) in enumerate(average_ratings_ceil.items()):
                row = i // n_cols
                col = i % n_cols
                # 计算每个星级的数量
                star_counts = (
                    df_filtered[df_filtered["condition"] == condition]["rating"]
                    .value_counts()
                    .apply(np.ceil)
                    .reindex(range(1, 11), fill_value=0)
                )
                # 绘制饼状图
                ax = axes[row, col] if n_cols > 1 else axes[row]
                ax.pie(
                    star_counts,
                    labels=star_counts.index,
                    autopct="%1.1f%%",
                    startangle=140,
                )
                ax.set_title(f"{condition} - Star Rating Distribution")
                ax.axis("equal")  # 使饼图为正圆形
            # 隐藏剩余的子图
            for j in range(i + 1, n_rows * n_cols):
                fig.delaxes(axes.flatten()[j])
            plt.tight_layout()
            plt.show()
        else:
            print("Error: 'condition' column not found in the DataFrame.")

    def RelationAnalysis(self):
        # 处理rating字段
        self.df["rating"] = self.df["rating"].apply(
            lambda x: (
                0 if 1 <= x <= 3 else 1 if 4 <= x <= 5 else 2 if 6 <= x <= 8 else 3
            )
        )
        self.df["effectiveness"] = self.df["effectiveness"].apply(
            self.map_effectiveness
        )
        self.df["sideEffects"] = self.df["sideEffects"].apply(self.map_side_effects)
        self.df["benefitsReview_len"] = self.df["benefitsReview"].apply(
            self.map_review_length
        )
        self.df["sideEffectsReview_len"] = self.df["sideEffectsReview"].apply(
            self.map_review_length
        )
        self.df["commentsReview_len"] = self.df["commentsReview"].apply(
            self.map_review_length
        )
        # 将分类变量转换为布尔型项集格式
        transactions = []
        for index, row in self.df.iterrows():
            transaction = [
                "rating_{}".format(row["rating"]),
                "effectiveness_{}".format(row["effectiveness"]),
                "sideEffects_{}".format(row["sideEffects"]),
                "benefitsReview_len_{}".format(row["benefitsReview_len"]),
                "sideEffectsReview_len_{}".format(row["sideEffectsReview_len"]),
                "commentsReview_len_{}".format(row["commentsReview_len"]),
            ]
            transactions.append(transaction)
        # 使用TransactionEncoder
        te = TransactionEncoder()
        te_ary = te.fit(transactions).transform(transactions)
        df_encoded = pd.DataFrame(te_ary, columns=te.columns_)
        # 使用apriori()算法找出频繁项集
        frequent_itemsets = apriori(df_encoded, min_support=0.01, use_colnames=True)
        # 使用association_rules()函数计算关联规则
        rules = association_rules(
            frequent_itemsets, metric="confidence", min_threshold=0.1, num_itemsets=10
        )
        # 将小数点后的位数限制为两位
        rules = rules.round(2)
        # 显示关联规则
        pd.set_option("display.max_rows", None)  # 设置显示所有行
        print(rules)
        # 如果你想要将结果保存到CSV文件
        rules.to_csv("../OutputData/association_rules.csv", index=False)


# 使用类
if __name__ == "__main__":
    analyzer = AssociationAnalyzer("drugLibTest_raw")
    analyzer.ConditionDevide()
    analyzer.StarDevide()
    analyzer.RelationAnalysis()
